{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qN8P0AnTnAhh"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MnUwFbCAKB2r"
      },
      "source": [
        "## Before we start\n",
        "To edit the colab notebook, please go to \"File\" -> \"Save a copy in Drive\" and make any edits on your copy.\n",
        "\n",
        "Before we start, please run the following to make sure that your environment is\n",
        "correctly setup. If you don't see a greeting, please refer to the\n",
        "[Installation](../install.md) guide for instructions. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "both",
        "id": "ZrGitA_KnRO0"
      },
      "outputs": [],
      "source": [
        "#@title Upgrade tensorflow_federated and load TensorBoard\n",
        "#@test {\"skip\": true}\n",
        "!pip install --quiet --upgrade tensorflow-federated-nightly\n",
        "!pip install --quiet --upgrade nest-asyncio\n",
        "\n",
        "import nest_asyncio\n",
        "nest_asyncio.apply()\n",
        "\n",
        "%load_ext tensorboard\n",
        "\n",
        "import sys\n",
        "\n",
        "if not sys.warnoptions:\n",
        "    import warnings\n",
        "    warnings.simplefilter(\"ignore\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "both",
        "id": "8BKyHkMxKHfV"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "import collections\n",
        "from matplotlib import pyplot as plt\n",
        "from IPython.display import display, HTML, IFrame\n",
        "\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import tensorflow_federated as tff\n",
        "\n",
        "tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)\n",
        "\n",
        "np.random.seed(0)\n",
        "\n",
        "def greetings():\n",
        "  display(HTML('<b><font size=\"6\" color=\"#ff00f4\">Greetings, virtual tutorial participants!</font></b>'))\n",
        "  return True\n",
        "l = tff.federated_computation(greetings)()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AftvNA5VMemJ"
      },
      "source": [
        "# TensorFlow Federated for Image Classification"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "coAumH42q9nz"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/federated/tutorials/federated_learning_for_image_classification\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/federated/blob/v0.14.0/docs/tutorials/federated_learning_for_image_classification.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/federated/blob/v0.14.0/docs/tutorials/federated_learning_for_image_classification.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Zs2LgZBOMt4M"
      },
      "source": [
        "Let's experiment with federated learning in simulation. In this tutorial, we use the classic MNIST training example to introduce the\n",
        "Federated Learning (FL) API layer of TFF, `tff.learning` - a set of\n",
        "higher-level interfaces that can be used to perform common types of federated\n",
        "learning tasks, such as federated training, against user-supplied models\n",
        "implemented in TensorFlow.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dN7n8RS7-rLR"
      },
      "source": [
        "# Tutorial Outline\n",
        "We'll be training a model to perform image classification using the classic MNIST dataset, with the neural net learning to classify digit from image. In this case, we'll be simulating federated learning, with the training data distributed on different devices.\n",
        "\n",
        "<p><b>Sections</b></p>\n",
        "\n",
        "\n",
        "1.   Load TFF Libraries.\n",
        "2.   Explore/preprocess federated EMNIST dataset.\n",
        "3.   Create a model.\n",
        "4.   Set up federated averaging process for training.\n",
        "5.   Analyze training metrics.\n",
        "6.   Set up federated evaluation computation.\n",
        "7.   Analyze evaluation metrics.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5Cyy2AWbLMKj"
      },
      "source": [
        "## Preparing the input data\n",
        "\n",
        "Let's start with the data. Federated learning requires a federated data set,\n",
        "i.e., a collection of data from multiple users. Federated data is typically\n",
        "non-[i.i.d.](https://en.wikipedia.org/wiki/Independent_and_identically_distributed_random_variables),\n",
        "which poses a unique set of challenges. Users typically have different distributions of data depending on usage patterns.\n",
        "\n",
        "In order to facilitate experimentation, we seeded the TFF repository with a few\n",
        "datasets.\n",
        "\n",
        "Here's how we can load our sample dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bP6WDENHSSZ9"
      },
      "outputs": [],
      "source": [
        "# Code for loading federated data from TFF repository\n",
        "emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yeX8BKgPfeFw"
      },
      "source": [
        "The data sets returned by `load_data()` are instances of\n",
        "`tff.simulation.datasets.ClientData`, an interface that allows you to enumerate the set\n",
        "of users, to construct a `tf.data.Dataset` that represents the data of a\n",
        "particular user, and to query the structure of individual elements.\n",
        "\n",
        "Let's explore the dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kN4-U5nJgKig"
      },
      "outputs": [],
      "source": [
        "len(emnist_train.client_ids)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZyCzIrSegT62"
      },
      "outputs": [],
      "source": [
        "# Let's look at the shape of our data\n",
        "example_dataset = emnist_train.create_tf_dataset_for_client(\n",
        "    emnist_train.client_ids[0])\n",
        "\n",
        "example_dataset.element_spec"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EsvSXGEMgd9G"
      },
      "outputs": [],
      "source": [
        "# Let's select an example dataset from one of our simulated clients\n",
        "example_dataset = emnist_train.create_tf_dataset_for_client(\n",
        "    emnist_train.client_ids[0])\n",
        "\n",
        "# Your code to get an example element from one client:\n",
        "example_element = next(iter(example_dataset))\n",
        "\n",
        "example_element['label'].numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OmLV0nfMg98V"
      },
      "outputs": [],
      "source": [
        "plt.imshow(example_element['pixels'].numpy(), cmap='gray', aspect='equal')\n",
        "plt.grid(False)\n",
        "_ = plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y-3KUf0kC8TN"
      },
      "source": [
        "**Exploring non-iid data**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "veWNAxdQfrgZ"
      },
      "outputs": [],
      "source": [
        "## Example MNIST digits for one client\n",
        "f = plt.figure(figsize=(20,4))\n",
        "j = 0\n",
        "\n",
        "for e in example_dataset.take(40):\n",
        "  plt.subplot(4, 10, j+1)\n",
        "  plt.imshow(e['pixels'].numpy(), cmap='gray', aspect='equal')\n",
        "  plt.axis('off')\n",
        "  j += 1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_XJRWDFWniik"
      },
      "outputs": [],
      "source": [
        "# Number of examples per layer for a sample of clients\n",
        "f = plt.figure(figsize=(12,7))\n",
        "f.suptitle(\"Label Counts for a Sample of Clients\")\n",
        "for i in range(6):\n",
        "  ds = emnist_train.create_tf_dataset_for_client(emnist_train.client_ids[i])\n",
        "  k = collections.defaultdict(list)\n",
        "  for e in ds:\n",
        "    k[e['label'].numpy()].append(e['label'].numpy())\n",
        "  plt.subplot(2, 3, i+1)\n",
        "  plt.title(\"Client {}\".format(i))\n",
        "  for j in range(10):\n",
        "    plt.hist(k[j], density=False, bins=[0,1,2,3,4,5,6,7,8,9,10])\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JOUI4zW9LQNH"
      },
      "outputs": [],
      "source": [
        "# Let's play around with the emnist_train dataset.\n",
        "# Let's explore the non-iid charateristic of the example data.\n",
        "\n",
        "for i in range(5):\n",
        "  ds = emnist_train.create_tf_dataset_for_client(emnist_train.client_ids[i])\n",
        "  k = collections.defaultdict(list)\n",
        "  for e in ds:\n",
        "    k[e['label'].numpy()].append(e['pixels'].numpy())\n",
        "  f = plt.figure(i, figsize=(12,5))\n",
        "  f.suptitle(\"Client #{}'s Mean Image Per Label\".format(i))\n",
        "  for j in range(10):\n",
        "    mn_img = np.mean(k[j],0)\n",
        "    plt.subplot(2, 5, j+1)\n",
        "    plt.imshow(mn_img.reshape((28,28)))#,cmap='gray') \n",
        "    plt.axis('off')\n",
        "\n",
        "# Each client has different mean images -- each client will be nudging the model\n",
        "# in their own directions."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lMd01egqy9we"
      },
      "source": [
        "### Preprocessing the data\n",
        "\n",
        "Since the data is already a `tf.data.Dataset`,  preprocessing can be accomplished using Dataset transformations. [See here](https://www.tensorflow.org/guide/data) for more detail on these transformations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cyG_BMraSuu_"
      },
      "outputs": [],
      "source": [
        "NUM_CLIENTS = 10\n",
        "NUM_EPOCHS = 5\n",
        "BATCH_SIZE = 20\n",
        "SHUFFLE_BUFFER = 100\n",
        "PREFETCH_BUFFER=10\n",
        "\n",
        "def preprocess(dataset):\n",
        "\n",
        "  def batch_format_fn(element):\n",
        "    \"\"\"Flatten a batch `pixels` and return the features as an `OrderedDict`.\"\"\"\n",
        "    return collections.OrderedDict(\n",
        "        x=tf.reshape(element['pixels'], [-1, 784]),\n",
        "        y=tf.reshape(element['label'], [-1, 1]))\n",
        "\n",
        "  return dataset.repeat(NUM_EPOCHS).shuffle(SHUFFLE_BUFFER).batch(\n",
        "      BATCH_SIZE).map(batch_format_fn).prefetch(PREFETCH_BUFFER)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m9LXykN_jlJw"
      },
      "source": [
        "Let's verify this worked."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VChB7LMQjkYz"
      },
      "outputs": [],
      "source": [
        "preprocessed_example_dataset = preprocess(example_dataset)\n",
        "\n",
        "sample_batch = tf.nest.map_structure(lambda x: x.numpy(),\n",
        "                                     next(iter(preprocessed_example_dataset)))\n",
        "\n",
        "sample_batch"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JGsMvRQt9Agl"
      },
      "source": [
        "Here's a simple helper function that will construct a list of datasets from the\n",
        "given set of users as an input to a round of training or evaluation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_PHMvHAI9xVc"
      },
      "outputs": [],
      "source": [
        "def make_federated_data(client_data, client_ids):\n",
        "  return [\n",
        "      preprocess(client_data.create_tf_dataset_for_client(x))\n",
        "      for x in client_ids\n",
        "  ]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0M9PfjOtAVqw"
      },
      "source": [
        "Now, how do we choose clients?"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GZ6NYHxB8xer"
      },
      "outputs": [],
      "source": [
        "sample_clients = emnist_train.client_ids[0:NUM_CLIENTS]\n",
        "\n",
        "# Your code to get the federated dataset here for the sampled clients:\n",
        "federated_train_data = make_federated_data(emnist_train, sample_clients)\n",
        "\n",
        "print('Number of client datasets: {l}'.format(l=len(federated_train_data)))\n",
        "print('First dataset: {d}'.format(d=federated_train_data[0]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HOxq4tbi9m8-"
      },
      "source": [
        "## Creating a model with Keras\n",
        "\n",
        "If you are using Keras, you likely already have code that constructs a Keras\n",
        "model. Here's an example of a simple model that will suffice for our needs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LYCsJGJFWbqt"
      },
      "outputs": [],
      "source": [
        "def create_keras_model():\n",
        "  return tf.keras.models.Sequential([\n",
        "      tf.keras.layers.InputLayer(input_shape=(784,)),\n",
        "      tf.keras.layers.Dense(10, kernel_initializer='zeros'),\n",
        "      tf.keras.layers.Softmax(),\n",
        "  ])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A0214iKjCTyX"
      },
      "source": [
        "**Centralized training with Keras**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g5XW_p4iLlJ2"
      },
      "outputs": [],
      "source": [
        "## Centralized training with keras ---------------------------------------------\n",
        "\n",
        "# This is separate from the TFF tutorial, and demonstrates how to train a\n",
        "# Keras model in a centralized fashion (contrasting training in a federated env)\n",
        "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()\n",
        "\n",
        "# Preprocess the data (these are NumPy arrays)\n",
        "x_train = x_train.reshape(60000, 784).astype(\"float32\") / 255\n",
        "\n",
        "y_train = y_train.astype(\"float32\")\n",
        "\n",
        "mod = create_keras_model()\n",
        "mod.compile(\n",
        "    optimizer=tf.keras.optimizers.RMSprop(),\n",
        "    loss=tf.keras.losses.SparseCategoricalCrossentropy(),\n",
        "    metrics=[tf.keras.metrics.SparseCategoricalAccuracy()]\n",
        ")\n",
        "h = mod.fit(\n",
        "    x_train,\n",
        "    y_train,\n",
        "    batch_size=64,\n",
        "    epochs=2\n",
        ")\n",
        "\n",
        "# ------------------------------------------------------------------------------"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l6pzsjoJQ2_D"
      },
      "source": [
        "**Federated training using a Keras model**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NHdraKFH4OU2"
      },
      "source": [
        "In order to use any model with TFF, it needs to be wrapped in an instance of the\n",
        "`tff.learning.Model` interface."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mA8cJoGE3Rh_"
      },
      "source": [
        "More keras metrics you can add are [found here](https://www.tensorflow.org/api_docs/python/tf/keras/metrics)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q3ynrxd53HzY"
      },
      "outputs": [],
      "source": [
        "def model_fn():\n",
        "  # We _must_ create a new model here, and _not_ capture it from an external\n",
        "  # scope. TFF will call this within different graph contexts.\n",
        "  keras_model = create_keras_model()\n",
        "  return tff.learning.from_keras_model(\n",
        "      keras_model,\n",
        "      input_spec=preprocessed_example_dataset.element_spec,\n",
        "      loss=tf.keras.losses.SparseCategoricalCrossentropy(),\n",
        "      metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XJ5E3O18_JZ6"
      },
      "source": [
        "## Training the model on federated data\n",
        "\n",
        "Now that we have a model wrapped as `tff.learning.Model` for use with TFF, we\n",
        "can let TFF construct a Federated Averaging algorithm by invoking the helper\n",
        "function `tff.learning.build_federated_averaging_process`, as follows."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sk6mjOfycX5N"
      },
      "outputs": [],
      "source": [
        "iterative_process = tff.learning.build_federated_averaging_process(\n",
        "    model_fn,\n",
        "    client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.02),\n",
        "    # Add server optimizer here!\n",
        "    server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.0))\n",
        "    "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f8FpvN2n67sm"
      },
      "source": [
        "What just happened? TFF has constructed a pair of *federated computations* and\n",
        "packaged them into a `tff.templates.IterativeProcess` in which these computations\n",
        "are available as a pair of properties `initialize` and `next`.\n",
        "\n",
        "An iterative process will usually be driven by a control loop like:\n",
        "```\n",
        "def initialize():\n",
        "  ...\n",
        "\n",
        "def next(state):\n",
        "  ...\n",
        "\n",
        "iterative_process = IterativeProcess(initialize, next)\n",
        "state = iterative_process.initialize()\n",
        "for round in range(num_rounds):\n",
        "  state = iterative_process.next(state)\n",
        "```\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v1gbHQ_7BiyT"
      },
      "source": [
        "Let's invoke the `initialize` computation to construct the server state."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6cagCWlZmcch"
      },
      "outputs": [],
      "source": [
        "state = iterative_process.initialize()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TjjxTx9e_rMd"
      },
      "source": [
        "The second of the pair of federated computations, `next`, represents a single\n",
        "round of Federated Averaging, which consists of pushing the server state\n",
        "(including the model parameters) to the clients, on-device training on their\n",
        "local data, collecting and averaging model updates, and producing a new updated\n",
        "model at the server.\n",
        "\n",
        "Let's run a single round of training and visualize the results. We can use the\n",
        "federated data we've already generated above for a sample of users."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "F3M_W9dDE6Tm"
      },
      "outputs": [],
      "source": [
        "# Run one single round of training.\n",
        "state, metrics = iterative_process.next(state, federated_train_data)\n",
        "print('round  1, metrics={}'.format(metrics['train']))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UmhReXt9G4A5"
      },
      "source": [
        "Let's run a few more rounds. As noted earlier, typically at this point you would\n",
        "pick a subset of your simulation data from a new randomly selected sample of\n",
        "users for each round in order to simulate a realistic deployment in which users\n",
        "continuously come and go, but in this interactive notebook, for the sake of\n",
        "demonstration we'll just reuse the same users, so that the system converges\n",
        "quickly."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qrJkQuCRJP9C"
      },
      "outputs": [],
      "source": [
        "NUM_ROUNDS = 11\n",
        "for round_num in range(2, NUM_ROUNDS):\n",
        "  state, metrics = iterative_process.next(state, federated_train_data)\n",
        "  print('round {:2d}, metrics={}'.format(round_num, metrics['train']))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "joHYzn9jcs0Y"
      },
      "source": [
        "Training loss is decreasing after each round of federated training, indicating\n",
        "the model is converging. There are some important caveats with these training\n",
        "metrics, however, see the section on *Evaluation* later in this tutorial."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ruSHJl1IjhNf"
      },
      "source": [
        "##Displaying model metrics in TensorBoard\n",
        "Next, let's visualize the metrics from these federated computations using Tensorboard.\n",
        "\n",
        "Let's start by creating the directory and the corresponding summary writer to write the metrics to.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "E3QUBK41lWDW"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "import os\n",
        "import shutil\n",
        "\n",
        "logdir = \"/tmp/logs/scalars/training/\"\n",
        "if os.path.exists(logdir):\n",
        "  shutil.rmtree(logdir)\n",
        "\n",
        "# Your code to create a summary writer:\n",
        "summary_writer = tf.summary.create_file_writer(logdir)\n",
        "\n",
        "state = iterative_process.initialize()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w-2aGxUlzS_J"
      },
      "source": [
        "Plot the relevant scalar metrics with the same summary writer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JZtr4_8lzN-V"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "with summary_writer.as_default():\n",
        "  for round_num in range(1, NUM_ROUNDS):\n",
        "    state, metrics = iterative_process.next(state, federated_train_data)\n",
        "    for name, value in metrics['train'].items():\n",
        "      tf.summary.scalar(name, value, step=round_num)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iUouyAHG0Mk8"
      },
      "source": [
        "Start TensorBoard with the root log directory specified above. It can take a few seconds for the data to load."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "urYYcmA9089p"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "%tensorboard --logdir /tmp/logs/scalars/ --port=0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jejrFEVP1EDs"
      },
      "source": [
        "In order to view evaluation metrics the same way, you can create a separate eval folder, like \"logs/scalars/eval\", to write to TensorBoard."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m7lz59lMJ0kj"
      },
      "source": [
        "## Evaluation\n",
        "To perform evaluation on federated data, you can construct another *federated\n",
        "computation* designed for just this purpose, using the\n",
        "`tff.learning.build_federated_evaluation` function, and passing in your model\n",
        "constructor as an argument."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nRiXyqnXM2VO"
      },
      "outputs": [],
      "source": [
        "# Construct federated evaluation computation here:\n",
        "evaluation = tff.learning.build_federated_evaluation(model_fn)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SpfgdNDoRjPy"
      },
      "source": [
        "Now, let's compile a test sample of federated data and rerun evaluation on the\n",
        "test data. The data will come from a different sample of users, but from a\n",
        "distinct held-out data set."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "in8vProVNc04"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "shuffled_ids = emnist_test.client_ids.copy()\n",
        "random.shuffle(shuffled_ids)\n",
        "sample_clients = shuffled_ids[0:NUM_CLIENTS]\n",
        "\n",
        "federated_test_data = make_federated_data(emnist_test, sample_clients)\n",
        "\n",
        "len(federated_test_data), federated_test_data[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ty-ZwfE0NJfV"
      },
      "outputs": [],
      "source": [
        "# Run evaluation on the test data here, using the federated model produced from \n",
        "# training:\n",
        "test_metrics = evaluation(state.model, federated_test_data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e5fGtIJYNqYH"
      },
      "outputs": [],
      "source": [
        "str(test_metrics)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "67vYxrDWzRcj"
      },
      "source": [
        "This concludes the tutorial. We encourage you to play with the\n",
        "parameters (e.g., batch sizes, number of users, epochs, learning rates, etc.), to modify the code above to simulate training on random samples of users in\n",
        "each round, and to explore the other tutorials we've developed."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4Zv28F7QLo8O"
      },
      "source": [
        "# Build your own FL algorithms\n",
        "\n",
        "In the previous tutorials, we learned how to set up model and data pipelines, and use these to perform federated training using the `tff.learning` API.\n",
        "\n",
        "Of course, this is only the tip of the iceberg when it comes to FL research. In this tutorial, we are going to discuss how to implement federated learning algorithms *without* deferring to the `tff.learning` API. We aim to accomplish the following:\n",
        "\n",
        "**Goals:**\n",
        "\n",
        "\n",
        "*   Understand the general structure of federated learning algorithms.\n",
        "*   Explore the *Federated Core* of TFF.\n",
        "*   Use the Federated Core to implement Federated Averaging directly.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hQ_N9XbULo8P"
      },
      "source": [
        "## Preparing the input data\n",
        "We first load and preprocess the EMNIST dataset included in TFF. We essentially use the same code as in the first tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-WdnFluLLo8P"
      },
      "outputs": [],
      "source": [
        "emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Blrh8zJgLo8R"
      },
      "outputs": [],
      "source": [
        "NUM_CLIENTS = 10\n",
        "BATCH_SIZE = 20\n",
        "\n",
        "def preprocess(dataset):\n",
        "\n",
        "  def batch_format_fn(element):\n",
        "    \"\"\"Flatten a batch of EMNIST data and return a (features, label) tuple.\"\"\"\n",
        "    return (tf.reshape(element['pixels'], [-1, 784]), \n",
        "            tf.reshape(element['label'], [-1, 1]))\n",
        "\n",
        "  return dataset.batch(BATCH_SIZE).map(batch_format_fn)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-vYM_IT7Lo8W"
      },
      "outputs": [],
      "source": [
        "client_ids = np.random.choice(emnist_train.client_ids, size=NUM_CLIENTS, replace=False)\n",
        "\n",
        "federated_train_data = [preprocess(emnist_train.create_tf_dataset_for_client(x))\n",
        "  for x in client_ids\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gNO_Y9j_Lo8X"
      },
      "source": [
        "## Preparing the model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LJ0I89ixz8yV"
      },
      "source": [
        "We use the same model as the first tutorial, which has a single hidden layer, followed by a softmax layer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Yfld4oFNLo8Y"
      },
      "outputs": [],
      "source": [
        "def create_keras_model():\n",
        "  return tf.keras.models.Sequential([\n",
        "      tf.keras.layers.InputLayer(input_shape=(784,)),\n",
        "      tf.keras.layers.Dense(10, kernel_initializer='zeros'),\n",
        "      tf.keras.layers.Softmax(),\n",
        "  ])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vLln0Q8G0Bky"
      },
      "source": [
        "We wrap this Keras model as a `tff.learning.Model`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SPwbipTNLo8a"
      },
      "outputs": [],
      "source": [
        "def model_fn():\n",
        "  keras_model = create_keras_model()\n",
        "  return tff.learning.from_keras_model(\n",
        "      keras_model,\n",
        "      input_spec=federated_train_data[0].element_spec,\n",
        "      loss=tf.keras.losses.SparseCategoricalCrossentropy(),\n",
        "      metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fPOWP2JjsfTk"
      },
      "source": [
        "# Cutomizing FL Algorithm\n",
        "\n",
        "While the `tff.learning` API encompasses many variants of Federated Averaging, there are many other algorithms that do not fit neatly into this framework. For example, you may want to add regularization, clipping, or more complicated algorithms such as [federated GAN training](https://github.com/google-research/federated/blob/master/gans). You may also be instead be interested in [federated analytics](https://ai.googleblog.com/2020/05/federated-analytics-collaborative-data.html)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "50N36Zz8qyY-"
      },
      "source": [
        "For these more advanced algorithms, we'll have to write our own custom FL algorithm.\n",
        "\n",
        "In general, FL algorithms have 4 main components:\n",
        "\n",
        "1. A server-to-client broadcast step.\n",
        "2. A local client update step.\n",
        "3. A client-to-server upload step.\n",
        "4. A server update step."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jH8s0GRdQt3b"
      },
      "source": [
        "In TFF, we generally represent federated algorithms as an `IterativeProcess`. This is simply a class that contains an `initialize_fn` and a `next_fn`. The `initialize_fn` will be used to initialize the server, and the `next_fn` will perform one communication round of federated averaging. Let's write a skeleton of what our iterative process for FedAvg should look like.\n",
        "\n",
        "First, we have an initialize function that simply creates a `tff.learning.Model`, and returns its trainable weights."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ylLpRa7T5DDh"
      },
      "outputs": [],
      "source": [
        "def initialize_fn():\n",
        "  model = model_fn()\n",
        "  return model.weights.trainable"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nb1-XAK8fB2A"
      },
      "source": [
        "This function looks good, but as we will see later, we will need to make a small modification to make it a TFF computation.\n",
        "\n",
        "We also want to sketch the `next_fn`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IeHN-XLZfMso"
      },
      "outputs": [],
      "source": [
        "def next_fn(server_weights, federated_dataset):\n",
        "  # Broadcast the server weights to the clients.\n",
        "  server_weights_at_client = broadcast(server_weights)\n",
        "\n",
        "  # Each client computes their updated weights.\n",
        "  client_weights = client_update(federated_dataset, server_weights_at_client)\n",
        "\n",
        "  # The server averages these updates.\n",
        "  mean_client_weights = mean(client_weights)\n",
        "\n",
        "  # The server updates its model.\n",
        "  server_weights = server_update(mean_client_weights)\n",
        "\n",
        "  return server_weights"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uWXvjXPWeujU"
      },
      "source": [
        "We'll focus on implementing these four components separately. We'll first focus on the parts that can be implemented in pure TensorFlow, namely the client and server update steps.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3yKS4VkALo8g"
      },
      "source": [
        "## TensorFlow Blocks "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bxpNYucgLo8g"
      },
      "source": [
        "### Client update\n",
        "\n",
        "We will use our `tff.learning.Model` to do client training in essentially the same way you would train a TF model. In particular, we will use `tf.GradientTape` to compute the gradient on batches of data, then apply these gradient using a `client_optimizer`.\n",
        "\n",
        "Note that each `tff.learning.Model` instance has a `weights` attribute with two sub-attributes:\n",
        "\n",
        "\n",
        "*   `trainable`: A list of the tensors corresponding to trainable layers.\n",
        "*   `non_trainable`: A list of the tensors corresponding to non-trainable layers.\n",
        "\n",
        "For our purposes, we will only use the trainable weights (as our model only has those!).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c5rHPKreLo8g"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def client_update(model, dataset, server_weights, client_optimizer):\n",
        "  \"\"\"Performs training (using the server model weights) on the client's dataset.\"\"\"\n",
        "  # Initialize the client model with the current server weights.\n",
        "  client_weights = model.weights.trainable\n",
        "  # Assign the server weights to the client model.\n",
        "  tf.nest.map_structure(lambda x, y: x.assign(y),\n",
        "                        client_weights, server_weights)\n",
        "\n",
        "  # Use the client_optimizer to update the local model.\n",
        "  for batch in dataset:\n",
        "    with tf.GradientTape() as tape:\n",
        "      # Compute a forward pass on the batch of data\n",
        "      outputs = model.forward_pass(batch)\n",
        "\n",
        "    # Compute the corresponding gradient\n",
        "    grads = tape.gradient(outputs.loss, client_weights)\n",
        "    grads_and_vars = zip(grads, client_weights)\n",
        "\n",
        "    # Apply the gradient using a client optimizer.\n",
        "    client_optimizer.apply_gradients(grads_and_vars)\n",
        "\n",
        "  return client_weights"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pP0D9XtoLo8i"
      },
      "source": [
        "### Server Update\n",
        "\n",
        "The server update will require even less effort. We will implement vanilla federated averaging, in which we simply replace the server model weights by the average of the client model weights. Again, we will only focus on the trainable weights."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rYxErLvHLo8i"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def server_update(model, mean_client_weights):\n",
        "  \"\"\"Updates the server model weights as the average of the client model weights.\"\"\"\n",
        "  model_weights = model.weights.trainable\n",
        "  # Assign the mean client weights to the server model.\n",
        "  tf.nest.map_structure(lambda x, y: x.assign(y),\n",
        "                        model_weights, mean_client_weights)\n",
        "  return model_weights"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ddCklfWlVr1U"
      },
      "source": [
        "Note that the code snippet above is clearly overkill, as we could simply return `mean_client_weights`. However, more advanced implementations of Federated Averaging could use `mean_client_weights` with more sophisticated techniques, such as momentum or adaptivity."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KuP9g6RFLo8k"
      },
      "source": [
        "So far, we've only written pure TensorFlow code. This is by design, as TFF allows you to use much of the TensorFlow code you're already familiar with. However, now we have to specify the *orchestration logic*, that is, the logic that dictates what the server broadcasts to the client, and what the client uploads to the server.\n",
        "\n",
        "This will require the \"Federated Core\" of TFF."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0CgFLVPgLo8l"
      },
      "source": [
        "# Introduction to the Federated Core\n",
        "\n",
        "The Federated Core (FC) is a set of lower-level interfaces that serve as the foundation for the `tff.learning` API. However, these interfaces are not limited to learning. In fact, they can be used for analytics and many other computations over distributed data.\n",
        "\n",
        "At a high-level, the federated core is a development environment that enables compactly expressed program logic to combine TensorFlow code with distributed communication operators (such as distributed sums and broadcasts). The goal is to give researchers and practitioners expliict control over the distributed communication in their systems, without requiring system implementation details (such as specifying point-to-point network message exchanges).\n",
        "\n",
        "One key point is that TFF is designed for privacy-preservation. Therefore, it allows explicit control over where data resides, to prevent unwanted accumulation of data at the centralized server location."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EYinjNqZLo8l"
      },
      "source": [
        "## Federated data\n",
        "\n",
        "Similar to \"Tensor\" concept in TensorFlow, which is one of the fundamental concepts, a key concept in TFF is \"federated data\", which refers to a collection of data items hosted across a group of devices in a distributed system (eg. client datasets, or the server model weights). We model the entire collection of data items across all devices as a single *federated value*.\n",
        "\n",
        "For example, suppose we have client devices that each have a float representing the temperature of a sensor. We could represent it as a *federated float* by"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7EJY0MHpLo8l"
      },
      "outputs": [],
      "source": [
        "federated_float_on_clients = tff.type_at_clients(tf.float32)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JSQAXD0FLo8n"
      },
      "source": [
        "Federated types are specified by a type `T` of its member constituents (eg. `tf.float32`) and a group `G` of devices. We will focus on the cases where `G` is either `tff.CLIENTS` or `tff.SERVER`. Such a federated type is represented as `{T}@G`, as shown below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6mlPgubJLo8n"
      },
      "outputs": [],
      "source": [
        "str(federated_float_on_clients)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pjAQytkeLo8o"
      },
      "source": [
        "Why do we care so much about placements? A key goal of TFF is to enable writing code that could be deployed on a real distributed system. This means that it is vital to reason about which subsets of devices execute which code, and where different pieces of data reside.\n",
        "\n",
        "TFF focuses on three things: *data*, where the data is *placed*, and how the data is being *transformed*. The first two are encapsulated in federated types, while the last is encapsulated in *federated computations*."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZLT2FmVMLo8p"
      },
      "source": [
        "## Federated computations"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-XwDC1vTLo8p"
      },
      "source": [
        "TFF is a strongly-typed functional programming environment whose basic units are *federated computations*. These are pieces of logic that accept federated values as input, and return federated values as output.\n",
        "\n",
        "For example, suppose we wanted to average the temperatures on our client sensors. We could define the following (using our federated float):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IfwXDNR1Lo8p"
      },
      "outputs": [],
      "source": [
        "@tff.federated_computation(tff.type_at_clients(tf.float32))\n",
        "def get_average_temperature(client_temperatures):\n",
        "  return tff.federated_mean(client_temperatures)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iSgs6Te5Lo8r"
      },
      "source": [
        "You might ask, how is this different from the `tf.function` decorator in TensorFlow? The key answer is that the code generated by `tff.federated_computation` is neither TensorFlow nor Python code; It is a specification of a distributed system in an internal platform-independent *glue language*.\n",
        "\n",
        "While this may sound complicated, you can think of TFF computations as functions with well-defined type signatures. These type signatures can be directly queried."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wAG1eDlULo8r"
      },
      "outputs": [],
      "source": [
        "str(get_average_temperature.type_signature)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TveOYFfuLo8s"
      },
      "source": [
        "This `tff.federated_computation` accepts arguments of federated type `{float32}@CLIENTS`, and returns values of federated type `{float32}@SERVER`. Federated computations may also go from server to client, from client to client, or from server to server. Federated computations can also be composed like normal functions, as long as their type signatures match up.\n",
        "\n",
        "To support development, TFF allows you to invoke a `tff.federated_computation` as a Python function. For example, we can call"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eFoqtuOTLo8t"
      },
      "outputs": [],
      "source": [
        "get_average_temperature([68.5, 70.3, 69.8])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZXn-yje9RJ6H"
      },
      "source": [
        "## Non-eager computations and TensorFlow"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nwyj8f3HLo8w"
      },
      "source": [
        "There are two key restrictions to be aware of. First, when the Python interpreter encounters a `tff.federated_computation` decorator, the function is traced once and serialized for future use. Therefore, TFF computations are fundamentally *non-eager*. This behavior is somewhat analogous to that of the [`tf.function`](https://www.tensorflow.org/api_docs/python/tf/function) decorator in TensorFlow.\n",
        "\n",
        "Second, a federated computation can only consist of federated operators (such as `tff.federated_mean`), they cannot contain TensorFlow operations. TensorFlow code must be confined to blocks decorated with `tff.tf_computation`. Most ordinary TensorFlow code can be directly decotrated, such as the following function that takes a number and adds `0.5` to it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "huz3mNmMLo8w"
      },
      "outputs": [],
      "source": [
        "@tff.tf_computation(tf.float32)\n",
        "def add_half(x):\n",
        "  return tf.add(x, 0.5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5ptjWALDLo8y"
      },
      "source": [
        "These also have type signatures, but *without placements*. For example, we can call"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xfAb0vG2Lo8y"
      },
      "outputs": [],
      "source": [
        "str(add_half.type_signature)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WNjwrNMjLo8z"
      },
      "source": [
        "Here we see an important difference between `tff.federated_computation` and `tff.tf_computation`. The former has explicit placements, while the latter does not.\n",
        "\n",
        "We can use `tff.tf_computation` blocks in federated computations by specifying placements. Let's create a function that adds half, but only to federated floats at the clients. We can do this by using `tff.federated_map`, which applies a given `tff.tf_computation`, while preserving the placement."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pG6nw3wiLo80"
      },
      "outputs": [],
      "source": [
        "@tff.federated_computation(tff.type_at_clients(tf.float32))\n",
        "def add_half_on_clients(x):\n",
        "  return tff.federated_map(add_half, x)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h4msKRKJLo81"
      },
      "source": [
        "This function is almost identical to `add_half`, except that it only accepts values with placement at `tff.CLIENTS`, and returns values with the same placement. We can see this in its type signature:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_C2hDiz0Lo82"
      },
      "outputs": [],
      "source": [
        "str(add_half_on_clients.type_signature)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3JxQ0DeiLo83"
      },
      "source": [
        "In summary:\n",
        "\n",
        "*   TFF operates on federated values.\n",
        "*   Each federated value has a *federated type*, with a *type* (eg. `tf.float32`) and a *placement* (eg. `tff.CLIENTS`).\n",
        "*   Federated values can be transformed using *federated computations*, which must be decorated with `tff.federated_computation` and a federated type signature.\n",
        "*   TensorFlow code must be contained in blocks with `tff.tf_computation` decorators. \n",
        "*   These blocks can then be incorporated into federated computations.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PvyFWox3Lo83"
      },
      "source": [
        "# Building your own FL Algorithm (Part 2)\n",
        "\n",
        "Now that we've peeked at the Federated Core, we can build our own federated learning algorithm. Remember that above, we defined an `initialize_fn` and `next_fn` for our algorithm. The `next_fn` will make use of the `client_update` and `server_update` we defined using pure TensorFlow code.\n",
        "\n",
        "However, in order to make our algorithm a federated computation, we will need both the `next_fn` and `initialize_fn` to be `tff.federated_computations`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CvY8fh1cLo84"
      },
      "source": [
        "## TensorFlow Federated blocks "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g0zNTO7LLo84"
      },
      "source": [
        "### Creating the initialization computation\n",
        "\n",
        "The initialize function will be quite simple: We will create a model using `model_fn`. However, remember that we must separate out our TensorFlow code using `tff.tf_computation`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jJY9xUBZLo84"
      },
      "outputs": [],
      "source": [
        "@tff.tf_computation\n",
        "def server_init():\n",
        "  model = model_fn()\n",
        "  return model.weights.trainable"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SGlv8LLgLo85"
      },
      "source": [
        "We can then pass this directly into a federated computation using `tff.federated_value`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m2hinzuRLo86"
      },
      "outputs": [],
      "source": [
        "@tff.federated_computation\n",
        "def initialize_fn():\n",
        "  return tff.federated_value(server_init(), tff.SERVER)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NFBghOgxLo88"
      },
      "source": [
        "### Creating the `next_fn`\n",
        "\n",
        "We now use our client and server update code to write the actual algorithm. We will first turn our `client_update` into a `tff.tf_computation` that accepts a client datasets and server weights, and outputs an updated client weights tensor.\n",
        "\n",
        "We will need the corresponding types to properly decorate our function. Luckily, the type of the server weights can be extracted directly from our model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ph_noHN2Lo88"
      },
      "outputs": [],
      "source": [
        "whimsy_model = model_fn()\n",
        "tf_dataset_type = tff.SequenceType(whimsy_model.input_spec)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WMPgpTaW66qx"
      },
      "source": [
        "Let's look at the dataset type signature. Remember that we took 28 by 28 images (with integer labels) and flattened them."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ju91izuz64wD"
      },
      "outputs": [],
      "source": [
        "str(tf_dataset_type)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kuS8d0BHLo8-"
      },
      "source": [
        "We can also extract the model weights type by using our `server_init` function above."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4yx6CExMLo8-"
      },
      "outputs": [],
      "source": [
        "model_weights_type = server_init.type_signature.result"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mS-Eh6Xj7J15"
      },
      "source": [
        "Examining the type signature, we'll be able to see the architecture of our model!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nPdhhM1O7IIL"
      },
      "outputs": [],
      "source": [
        "str(model_weights_type)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g1U1wTGRLo8_"
      },
      "source": [
        "We can now create our `tff.tf_computation` for the client update."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q0W05pMWLo9A"
      },
      "outputs": [],
      "source": [
        "@tff.tf_computation(tf_dataset_type, model_weights_type)\n",
        "def client_update_fn(tf_dataset, server_weights):\n",
        "  model = model_fn()\n",
        "  client_optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)\n",
        "  return client_update(model, tf_dataset, server_weights, client_optimizer)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uP5quaAuLo9B"
      },
      "source": [
        "The `tff.tf_computation` version of the server update can be defined in a similar way, using types we've already extracted."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "F4WvQtVzLo9B"
      },
      "outputs": [],
      "source": [
        "@tff.tf_computation(model_weights_type)\n",
        "def server_update_fn(mean_client_weights):\n",
        "  model = model_fn()\n",
        "  return server_update(model, mean_client_weights)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SImhLbu4Lo9D"
      },
      "source": [
        "Last, but not least, we need to create the `tff.federated_computation` that brings this all together. This function will accept two *federated values*, one corresponding to the server weights (with placement `tff.SERVER`), and the other corresponding to the client datasets (with placement `tff.CLIENTS`).\n",
        "\n",
        "Note that both these types were defined above! We simply need to give them the proper placement using `tff.type_at_{server/clients}``."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ekPsA8AsLo9D"
      },
      "outputs": [],
      "source": [
        "federated_server_type = tff.type_at_server(model_weights_type)\n",
        "federated_dataset_type = tff.type_at_clients(tf_dataset_type)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7FXAX7vGLo9G"
      },
      "source": [
        "Remember the 4 elements of an FL algorithm?\n",
        "\n",
        "1. A server-to-client broadcast step.\n",
        "2. A local client update step.\n",
        "3. A client-to-server upload step.\n",
        "4. A server update step.\n",
        "\n",
        "Now that we've built up the above, each part can be compactly represented as a single line of TFF code. This simplicity is why we had to take extra care to specify things such as federated types!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Epc7MwfELo9G"
      },
      "outputs": [],
      "source": [
        "@tff.federated_computation(federated_server_type, federated_dataset_type)\n",
        "def next_fn(server_weights, federated_dataset):\n",
        "  # Broadcast the server weights to the clients.\n",
        "  server_weights_at_client = tff.federated_broadcast(server_weights)\n",
        "\n",
        "  # Each client computes their updated weights.\n",
        "  client_weights = tff.federated_map(\n",
        "      client_update_fn, (federated_dataset, server_weights_at_client))\n",
        "  \n",
        "  # The server averages these updates.\n",
        "  mean_client_weights = tff.federated_mean(client_weights)\n",
        "\n",
        "  # The server updates its model.\n",
        "  server_weights = tff.federated_map(server_update_fn, mean_client_weights)\n",
        "\n",
        "  return server_weights"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kWomG3TtLo9I"
      },
      "source": [
        "We now have a `tff.federated_computation` for both the algorithm initialization, and for running one step of the algorithm. To finish our algorithm, we pass these into `tff.templates.IterativeProcess`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GxdWgEddLo9I"
      },
      "outputs": [],
      "source": [
        "federated_algorithm = tff.templates.IterativeProcess(\n",
        "    initialize_fn=initialize_fn,\n",
        "    next_fn=next_fn\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7Z__9k-Dc1I3"
      },
      "source": [
        "Let's look at the *type signature* of the `initialize` and `next` functions of our iterative process."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kyRLXDj-Lo9J"
      },
      "outputs": [],
      "source": [
        "str(federated_algorithm.initialize.type_signature)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UyyEi5Kec90_"
      },
      "source": [
        "This reflects the fact that `federated_algorithm.initialize` is a no-arg function that returns a single-layer model (with a 784-by-10 weight matrix, and 10 bias units)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mx6yuIKtLo9M"
      },
      "outputs": [],
      "source": [
        "str(federated_algorithm.next.type_signature)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "efpdHodmdU_6"
      },
      "source": [
        "Here, we see that `federated_algorithm.next` accepts a server model and client data, and returns an updated server model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4UYZ3qeMLo9N"
      },
      "source": [
        "## Evaluating the algorithm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jwd9Gs0ULo9O"
      },
      "source": [
        "Let's run a few rounds, and see how the loss changes. First, we will define an evaluation function using the *centralized* approach discussed in the second tutorial.\n",
        "\n",
        "We first create a centralized evaluation dataset, and then apply the same preprocessing we used for the training data.\n",
        "\n",
        "Note that we only `take` the first 1000 elements for reasons of computational efficiency, but typically we'd use the entire test dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EdNgYoIwLo9P"
      },
      "outputs": [],
      "source": [
        "central_emnist_test = emnist_test.create_tf_dataset_from_all_clients().take(1000)\n",
        "central_emnist_test = preprocess(central_emnist_test)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7R50NZ35dphE"
      },
      "source": [
        "Next, we write a function that accepts a server state, and uses Keras to evaluate on the test dataset. If you're familiar with `tf.Keras`, this will all look familiar, though note the use of `set_weights`!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I5UEX4EWLo9Q"
      },
      "outputs": [],
      "source": [
        "def evaluate(server_state):\n",
        "  keras_model = create_keras_model()\n",
        "  keras_model.compile(\n",
        "      loss=tf.keras.losses.SparseCategoricalCrossentropy(),\n",
        "      metrics=[tf.keras.metrics.SparseCategoricalAccuracy()]  \n",
        "  )\n",
        "  keras_model.set_weights(server_state)\n",
        "  keras_model.evaluate(central_emnist_test)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hygoBACkLo9S"
      },
      "source": [
        "Now, let's initialize our algorithm and evaluate on the test set."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JC6zMDTTLo9S"
      },
      "outputs": [],
      "source": [
        "server_state = federated_algorithm.initialize()\n",
        "evaluate(server_state)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2knqix2cLo9U"
      },
      "source": [
        "Let's train for a few rounds and see if anything changes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v1zBlzFILo9U"
      },
      "outputs": [],
      "source": [
        "for round in range(15):\n",
        "  server_state = federated_algorithm.next(server_state, federated_train_data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "q91Vjyc_jumU"
      },
      "outputs": [],
      "source": [
        "evaluate(server_state)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XM34ammUW-T3"
      },
      "source": [
        "We see a slight decrease in the loss function. While the jump is small, note that we've only performed 10 training rounds, and on a small subset of clients. To see better results, we may have to do hundreds if not thousands of rounds."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o13H5dDFXRFn"
      },
      "source": [
        "## Modifying our algorithm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qt4jVD21XTL-"
      },
      "source": [
        "At this point, let's stop and think about what we've accomplished. We've implemented Federated Averaging directly by combining pure TensorFlow code (for the client and server updates) with federated computations from the Federated Core of TFF.\n",
        "\n",
        "To perform more sophisticted learning, we can simply alter what we have above. In particular, by editing the pure TF code above, we can change how the client performs training, or how the server updates its model.\n",
        "\n",
        "**Challenge:** Add [gradient clipping](https://towardsdatascience.com/what-is-gradient-clipping-b8e815cdfb48) to the `client_update` function.\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "openmined_conference_2020.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
